import asyncio

import sqlmodel
from sqlalchemy.exc import OperationalError
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlmodel import SQLModel, Session
import timescaledb
import motor.motor_asyncio

from .config import DATABASE_URL, DB_TIMEZONE, MONGO_URL

if DATABASE_URL == "":
    raise NotImplementedError("DATABASE_URL not set")
if MONGO_URL == "":
    raise NotImplementedError("MONGO_URL not set")

# 异步数据库引擎
sql_engine = create_async_engine(DATABASE_URL, echo=True)

# 创建异步会话工厂
AsyncSessionLocal = sessionmaker(sql_engine, class_=AsyncSession, expire_on_commit=False)
# MongoDB
mongo_client = motor.motor_asyncio.AsyncIOMotorClient(MONGO_URL)
mongo_db = mongo_client["myTrader"]


# 异步初始化数据库
async def init_sqldb(retries: int = 10, delay: int = 3):
    for attempt in range(retries):
        try:
            print(f"[DB INIT] Trying to connect to TimescaleDB (Attempt {attempt + 1}/{retries})")
            async with sql_engine.begin() as conn:
                # 创建 SQLModel 表
                await conn.run_sync(SQLModel.metadata.create_all)

            print("[DB INIT] Creating hypertables...")
            async with sql_engine.begin() as conn:
                await conn.run_sync(timescaledb.metadata.create_all)

            print("[DB INIT] ✅ TimescaleDB initialization completed.")
            return
        except OperationalError as e:
            print(f"[DB INIT] ❌ Connection failed: {e}")
            if attempt < retries - 1:
                print(f"[DB INIT] ⏳ Retrying in {delay} seconds...")
                await asyncio.sleep(delay)
            else:
                print("[DB INIT] ❌ All retries failed.")
                raise RuntimeError("Could not connect to database after retries.") from e


# 按 collection 获取的依赖函数
def get_collection(name: str):
    return mongo_db[name]